# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import json
import math
import sys
from typing import Iterable, Optional

import cv2
import torchvision
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
# from pytorch_grad_cam import GradCAM
# from pytorch_grad_cam.utils.image import show_cam_on_image
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report, confusion_matrix, roc_curve, auc, precision_recall_curve, average_precision_score
import torch
import numpy as np
from timm.data import Mixup
from timm.utils import accuracy
import PIL.Image
from PIL import Image
from torchcam.utils import overlay_mask
from torchvision.transforms.functional import to_pil_image

import util.misc as misc
import util.lr_sched as lr_sched
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import torchvision.transforms as transforms
# from torchcam.utils import overlay_mask

import os

# -----------------------------------------------
# 用做评估模型
#
# -----------------------------------------------


def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
                    mixup_fn: Optional[Mixup] = None, log_writer=None,
                    args=None):
    # 设置模型为训练模式
    model.train(True)

    # 初始化 MetricLogger，用于记录训练过程中各项指标
    metric_logger = misc.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))  # 记录学习率
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 20  # 每 20 个 iteration 打印一次日志

    # 获取每次梯度累积的迭代次数
    accum_iter = args.accum_iter

    # 清空优化器梯度
    optimizer.zero_grad()

    # 如果存在 log_writer，打印日志目录
    if log_writer is not None:
        print('log_dir: {}'.format(log_writer.log_dir))

    # 遍历数据加载器中的每个批次
    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):

        # 根据当前迭代步数调整学习率
        if data_iter_step % accum_iter == 0:
            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)

        # 将输入数据和目标标签移动到指定设备（GPU 或 CPU）
        samples = samples.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        # 如果启用了 mixup 数据增强，则执行混合操作
        if mixup_fn is not None:
            samples, targets = mixup_fn(samples, targets)

        # 使用自动混合精度（autocast）进行前向传播和计算损失
        with torch.cuda.amp.autocast():
            outputs = model(samples)  # 模型前向传播
            loss = criterion(outputs, targets)  # 计算损失

        loss_value = loss.item()  # 获取损失值

        # 检查损失值是否有效
        if not math.isfinite(loss_value):
            print("Loss is {}, stopping training".format(loss_value))  # 如果损失值无穷大或 NaN，则停止训练
            sys.exit(1)

        # 如果进行梯度累积，按累积次数对损失进行缩放
        loss /= accum_iter

        # 使用损失缩放器来更新梯度
        loss_scaler(loss, optimizer, clip_grad=max_norm,
                    parameters=model.parameters(), create_graph=False,
                    update_grad=(data_iter_step + 1) % accum_iter == 0)

        # 每经过 accum_iter 次迭代，清空优化器梯度
        if (data_iter_step + 1) % accum_iter == 0:
            optimizer.zero_grad()

        # 同步 GPU 操作，确保所有计算完成
        torch.cuda.synchronize()

        # 更新指标（损失）
        metric_logger.update(loss=loss_value)

        # 计算当前迭代的学习率范围（最小学习率和最大学习率）
        min_lr = 10.  # 初始化最小学习率
        max_lr = 0.  # 初始化最大学习率
        for group in optimizer.param_groups:
            min_lr = min(min_lr, group["lr"])  # 更新最小学习率
            max_lr = max(max_lr, group["lr"])  # 更新最大学习率

        # 更新学习率指标
        metric_logger.update(lr=max_lr)

        # 将当前的损失值进行全局平均，确保多个进程下的结果一致
        loss_value_reduce = misc.all_reduce_mean(loss_value)

        # 如果存在 log_writer，记录损失和学习率到 TensorBoard
        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
            # 使用 epoch_1000x 作为 x 轴，避免因为批大小变化导致的曲线不一致
            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
            log_writer.add_scalar('loss', loss_value_reduce, epoch_1000x)  # 记录损失值
            log_writer.add_scalar('lr', max_lr, epoch_1000x)  # 记录学习率

    # 汇总所有进程的统计数据
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)

    # 返回统计信息的平均值
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}


# 输出分类指标


@torch.no_grad()  # 禁用梯度计算，节省内存和计算资源
def evaluate(data_loader, model, device):
    criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失函数，用于分类任务

    output_dir = "./finetune_dir/opt/"  # 设置输出目录，保存结果
    os.makedirs(output_dir, exist_ok=True)  # 创建输出目录，如果目录已存在则不做任何操作

    # 初始化 MetricLogger，用于记录训练和评估过程中的各项指标
    metric_logger = misc.MetricLogger(delimiter="  ")
    header = 'Test:'  # 设置评估阶段的日志前缀

    model.eval()  # 将模型设置为评估模式（关闭 Dropout 等）

    # 初始化保存预测结果、真实标签和输出的列表
    predictions = []
    targets = []
    outputs = []
    res_preds = []

    # 遍历验证集数据
    for batch in metric_logger.log_every(data_loader, 10, header):  # 每10个batch输出一次日志
        images = batch[0]  # 获取当前批次的输入图像
        target = batch[-1]  # 获取当前批次的标签
        images = images.to(device, non_blocking=True)  # 将数据移到GPU
        target = target.to(device, non_blocking=True)  # 将标签移到GPU

        # 使用自动混合精度（autocast）来减少内存消耗并加速计算
        with torch.cuda.amp.autocast():
            output = model(images)  # 前向传播
            loss = criterion(output, target)  # 计算损失

            # 计算预测的类别概率值
            probabilities = torch.softmax(output, dim=1)
            predictions.extend(probabilities.cpu().numpy())  # 存储所有类的概率值
            targets.extend(target.cpu().numpy())  # 存储真实标签
            outputs.extend(images.cpu().numpy())  # 存储输入图像

        # 计算 top-1 和 top-5 准确率
        acc1, acc5 = accuracy(output, target, topk=(1, 2))

        batch_size = images.shape[0]  # 获取当前批次的大小
        metric_logger.update(loss=loss.item())  # 更新损失
        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)  # 更新 top-1 准确率
        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)  # 更新 top-5 准确率

    # 同步各进程的数据
    metric_logger.synchronize_between_processes()

    # 打印最终的评估结果：Top-1 准确率、Top-5 准确率、损失
    print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
          .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))

    # 将预测结果和真实标签转换为 NumPy 数组
    predictions = np.array(predictions)
    targets = np.array(targets)

    # 计算混淆矩阵
    pred_labels = np.argmax(predictions, axis=1)  # 获取预测的类别（概率值最大的位置）
    confusion_mat = confusion_matrix(targets, pred_labels)  # 计算混淆矩阵
    print(confusion_mat)  # 打印混淆矩阵

    # 计算 Precision、Recall、F1 Score 等分类评估指标
    classification_rep = classification_report(targets, pred_labels, digits=4, zero_division=0)
    print(classification_rep)  # 打印分类报告

    # 将分类报告写入文件
    with open(os.path.join(output_dir, "classification_report.txt"), 'w') as f:
        f.write(classification_rep)

    # 计算并绘制 Precision-Recall 曲线
    for i in range(3):  # 对于每个类别
        y_true = (targets == i).astype(int)  # 将真实标签转换为与当前类别对应的布尔数组（0或1）
        precision, recall, thresholds = precision_recall_curve(y_true, predictions[:, i])  # 计算精确度-召回率曲线
        average_precision = average_precision_score(y_true, predictions[:, i])  # 计算平均精度

        plt.clf()  # 清除当前图形
        plt.figure("P-R Curve")
        plt.title(f'Precision/Recall Curve (Class {i})')
        plt.xlabel('Recall')
        plt.ylabel('Precision')
        plt.plot(recall, precision, label='PR curve (area = {0:0.3f})'.format(average_precision))  # 绘制精确度-召回率曲线
        plt.legend(loc='best')  # 设置图例
        plt.savefig(os.path.join(output_dir, f"pr_curve_class_{i}.png"))  # 保存为图片

    # 绘制 ROC 曲线
    for i in range(3):
        y_true = (targets == i).astype(int)  # 将真实标签转换为与当前类别对应的布尔数组（0或1）
        fpr, tpr, _ = roc_curve(y_true, predictions[:, i])  # 计算假阳性率和真阳性率
        roc_auc = auc(fpr, tpr)  # 计算 ROC 曲线下面积

        plt.figure()  # 创建新图形
        lw = 2  # 线宽
        plt.plot(fpr, tpr, color='cornflowerblue', lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)  # 绘制 ROC 曲线
        plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')  # 绘制对角线
        plt.xlim([0.0, 1.0])  # 设置 x 轴范围
        plt.ylim([0.0, 1.05])  # 设置 y 轴范围
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title(f'ROC curve (Class {i})')  # 设置标题
        plt.legend(loc="lower right")  # 设置图例位置
        plt.savefig(os.path.join(output_dir, f"roc_class_{i}.png"))  # 保存为图片

    # 绘制混淆矩阵的热力图
    class_names = data_loader.dataset.classes  # 获取类别名称
    print(class_names)  # 打印类别名称
    labels = [class_names[0], class_names[1], class_names[2]]  # 只选择前3类进行绘制
    plt.figure(figsize=(8, 6))  # 设置图形大小
    sns.set(font_scale=1.2)  # 设置字体大小
    sns.heatmap(confusion_mat, annot=True, fmt="d", cmap="Blues", xticklabels=labels, yticklabels=labels)  # 绘制混淆矩阵热力图
    plt.xlabel("Predicted label")
    plt.ylabel("True label")
    plt.title("Confusion Matrix")
    plt.savefig(os.path.join(output_dir, "cm.jpg"))  # 保存混淆矩阵图像

    # 保存模型在验证集上的预测结果
    res_preds.extend(predictions)  # 将所有预测结果存储在 res_preds 列表中
    res_preds2 = pd.DataFrame(res_preds)  # 转换为 Pandas DataFrame 格式
    resppath = os.path.join(output_dir, "resps.txt")  # 设置保存路径
    res_preds2.to_csv(resppath, header=None, sep=' ', index=False)  # 将预测结果保存为 CSV 文件

    # 返回评估指标的平均值
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}



# 定义图像变换函数，调整特征图的形状以适应Grad-CAM可视化
def reshape_transform(tensor, height=14, width=14):
    # 去掉cls token并调整tensor形状
    result = tensor[:, 1:, :].reshape(tensor.size(0), height, width, tensor.size(2))

    # 将通道维度放到最后，并将空间维度调换位置
    result = result.transpose(2, 3).transpose(1, 2)
    return result

# 使用GradCAM进行可视化
def gt_cam(model, target_layers, inputs_tensor, img):
    # 检查是否可以使用GPU加速
    use_cuda = torch.cuda.is_available()
    if use_cuda:
        model = model.cuda()  # 将模型移到GPU

    # 创建GradCAM对象
    cam = GradCAM(model=model,
                  target_layers=[model.blocks[-1].norm1],  # 选择目标层，这里选择最后一层的规范化层
                  use_cuda=use_cuda,
                  reshape_transform=reshape_transform)  # 使用自定义的reshape_transform函数

    # 获取Grad-CAM图像
    grayscale_cam = cam(input_tensor=inputs_tensor)
    grayscale_cam = grayscale_cam[0, :]  # 选择第一个图像的Grad-CAM图像

    # 在原始图像上叠加Grad-CAM图像
    visualization = show_cam_on_image(img, grayscale_cam, use_rgb=True)
    # 返回可视化后的图像
    return visualization

# 实现传统的Grad-CAM方法
def grad_cam(model, input, target, layer_name):
    # 将模型设置为评估模式
    model.eval()

    # 获取目标层的特征图和梯度
    features, grads = None, None
    for name, module in model.named_modules():
        if name == layer_name:
            features = module.features  # 获取目标层的特征图
            grads = module.grads  # 获取目标层的梯度
            break

    # 前向传播并计算损失
    output = model(input)
    loss = torch.nn.CrossEntropyLoss()(output, target)
    loss.backward()  # 反向传播计算梯度

    # 计算梯度的权重
    weights = torch.mean(grads, dim=(2, 3))  # 按空间维度对梯度求均值
    weights = torch.relu(weights)  # 使用ReLU激活

    # 计算特征图加权和
    cam = torch.sum(weights[:, :, None, None] * features, dim=1)
    cam = torch.relu(cam)  # 使用ReLU激活

    # 对结果进行归一化
    cam = cam - torch.min(cam)  # 将最小值归零
    cam = cam / torch.max(cam)  # 将最大值归一化为1

    return cam

# 获取模型中最后一层卷积层
def get_last_conv(m):
    """
    获取一个模块中的最后一个卷积层
    """
    # 获取所有卷积层
    convs = filter(lambda k: isinstance(k, torch.nn.Conv2d), m.modules())
    return list(convs)[-1]  # 返回最后一个卷积层

# 对模型进行GT（Grad-CAM）可视化
def gt_vis(data_loader, model, device):
    # 切换到评估模式
    metric_logger = misc.MetricLogger(delimiter="  ")
    header = 'Test:'
    model.eval()  # 将模型设置为评估模式

    predictions = []  # 存储模型预测的结果
    targets = []  # 存储真实标签
    outputs = []  # 存储模型输出
    res_preds = []  # 存储所有的预测结果

    # 遍历数据加载器中的所有批次
    for batch in metric_logger.log_every(data_loader, 10, header):
        images = batch[0]  # 获取输入图像
        target = batch[-1]  # 获取真实标签
        images = images.to(device, non_blocking=True)  # 将图像移到设备上
        target = target.to(device, non_blocking=True)  # 将标签移到设备上

        # 模型进行前向预测
        output = model(images)
        prediction = output.argmax(1)  # 获取预测的类别标签

        # 记录预测结果和真实标签
        predictions.extend(prediction.cpu().numpy())
        targets.extend(target.cpu().numpy())
        outputs.extend(output.detach().cpu().numpy())  # 存储模型输出结果

        # 判断每个预测结果的类别与真实标签，分别保存不同情况的图像
        for w in range(len(target)):
            true_label = target[w].item()  # 获取真实标签
            pred_label = prediction[w].item()  # 获取预测标签
            origin_img = images[w].permute(1, 2, 0).cpu().numpy()  # 将图像从Tensor转为NumPy数组
            origin_img = (origin_img * 255).astype(np.uint8)  # 将像素值恢复到0-255范围并转换为uint8

            # 保存图像文件
            if true_label == 1:
                if true_label == pred_label:
                    cv2.imwrite(
                        f"./finetune_dir/visual/TP/original_batch_{w}_NO_true.png",  # TP：True Positive
                        cv2.cvtColor(origin_img, cv2.COLOR_RGB2BGR)
                    )
                else:
                    cv2.imwrite(
                        f"./finetune_dir/visual/TN/original_batch_{w}_NO.png",  # TN：True Negative
                        cv2.cvtColor(origin_img, cv2.COLOR_RGB2BGR)
                    )
            if true_label == 0:
                if true_label == pred_label:
                    cv2.imwrite(
                        f"./finetune_dir/visual/FP/original_batch_{w}_NO.png",  # FP：False Positive
                        cv2.cvtColor(origin_img, cv2.COLOR_RGB2BGR)
                    )
                else:
                    cv2.imwrite(
                        f"./finetune_dir/visual/FN/original_batch_{w}_NO.png",  # FN：False Negative
                        cv2.cvtColor(origin_img, cv2.COLOR_RGB2BGR)
                    )






